#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
指标预处理脚本

功能：
1. 遍历所有场景，使用MetricTool预计算所有组件的指标异常检测结果
2. 将结果保存到各场景的 metric_cache.json 文件中
3. 支持断点续传、并发处理、进度显示

使用方法：
    python scripts/preprocess_metrics.py --data-root /path/to/data --workers 4
    
    参数说明：
        --data-root: 数据集根目录（默认：data/envs/aiops-22）
        --workers: 并发线程数（默认：4）
        --force-refresh: 强制重新处理已有缓存的场景
        --scenarios: 仅处理指定场景（逗号分隔，如 "scene1,scene2"）
"""

import argparse
import sys
import os
import json
import logging
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Optional, Tuple
from datetime import datetime

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 导入项目模块
from ragen.env.aiops.tools.metric_tool import MetricTool
from ragen.env.aiops.metric_cache_loader import MetricCacheLoader

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('preprocess_metrics.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class MetricPreprocessor:
    """
    指标预处理器
    
    负责批量预处理所有场景的指标数据
    """
    
    def __init__(self, 
                 data_root: str,
                 workers: int = 4,
                 force_refresh: bool = False,
                 pre_fault_steps: int = 2,
                 contamination: float = 0.05):
        """
        初始化预处理器
        
        Args:
            data_root: 数据集根目录
            workers: 并发线程数
            force_refresh: 是否强制刷新已有缓存
            pre_fault_steps: MetricTool 参数
            contamination: MetricTool 参数
        """
        self.data_root = os.path.abspath(data_root)
        self.samples_json_path = os.path.join(self.data_root, 'samples.json')
        self.workers = workers
        self.force_refresh = force_refresh
        self.pre_fault_steps = pre_fault_steps
        self.contamination = contamination
        
        self.cache_loader = MetricCacheLoader()
        
        # 统计信息
        self.stats = {
            'total_scenarios': 0,
            'processed_scenarios': 0,
            'skipped_scenarios': 0,
            'failed_scenarios': 0,
            'total_components': 0,
            'start_time': datetime.now()
        }
        
        logger.info(f"预处理器初始化完成:")
        logger.info(f"  数据根目录: {self.data_root}")
        logger.info(f"  并发线程数: {self.workers}")
        logger.info(f"  强制刷新: {self.force_refresh}")
        logger.info(f"  pre_fault_steps: {self.pre_fault_steps}")
        logger.info(f"  contamination: {self.contamination}")
    
    def load_scenarios(self, specific_scenarios: Optional[List[str]] = None) -> List[Dict]:
        """
        加载场景列表
        
        Args:
            specific_scenarios: 指定要处理的场景ID列表（None表示处理全部）
            
        Returns:
            场景信息列表
        """
        try:
            with open(self.samples_json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            scenarios = data.get('scenarios', [])
            
            if specific_scenarios:
                # 过滤指定场景
                scenarios = [s for s in scenarios if s['scenario_id'] in specific_scenarios]
                logger.info(f"加载指定场景: {len(scenarios)} 个")
            else:
                logger.info(f"加载全部场景: {len(scenarios)} 个")
            
            self.stats['total_scenarios'] = len(scenarios)
            return scenarios
            
        except Exception as e:
            logger.error(f"加载场景列表失败: {e}")
            return []
    
    def process_scenario(self, scenario_info: Dict) -> Tuple[bool, str, int]:
        """
        处理单个场景
        
        Args:
            scenario_info: 场景信息字典
            
        Returns:
            (是否成功, 消息, 处理的组件数)
        """
        scenario_id = scenario_info['scenario_id']
        folder_name = scenario_info['folder_name']
        scenario_folder = os.path.join(self.data_root, folder_name)
        
        try:
            # 检查是否已有缓存
            if not self.force_refresh and self.cache_loader.cache_exists(scenario_folder):
                logger.info(f"⏭️  场景 {scenario_id} 已有缓存，跳过")
                self.stats['skipped_scenarios'] += 1
                return True, "已有缓存，跳过", 0
            
            # 创建 MetricTool 实例
            logger.info(f"🔄 开始处理场景: {scenario_id}")
            
            metric_tool = MetricTool(
                scenario_id=scenario_id,
                samples_json_path=self.samples_json_path,
                pre_fault_steps=self.pre_fault_steps,
                contamination=self.contamination,
                enable_cache=False,  # 禁用内部缓存，避免冲突
                max_workers=self.workers
            )
            
            # 获取所有可用的组件ID
            cmdb_ids = metric_tool.get_available_cmdb_ids()
            
            if not cmdb_ids:
                logger.warning(f"⚠️  场景 {scenario_id} 没有可用组件")
                return True, "无可用组件", 0
            
            logger.info(f"📊 场景 {scenario_id}: 发现 {len(cmdb_ids)} 个组件")
            
            # 批量检测异常
            def progress_callback(completed, total, cmdb_id):
                """进度回调"""
                logger.info(f"  [{scenario_id}] 进度: {completed}/{total} - {cmdb_id}")
            
            results = metric_tool.batch_detect_anomalies(
                cmdb_ids=cmdb_ids,
                train_on_normal_only=True,
                use_cache=False,
                progress_callback=progress_callback
            )
            
            # 保存到缓存文件
            metadata = {
                'pre_fault_steps': self.pre_fault_steps,
                'contamination': self.contamination,
                'tool_version': '1.0.0',
                'anomaly_component': scenario_info.get('anomaly_component', 'unknown'),
                'anomaly_reason': scenario_info.get('anomaly_reason', 'unknown')
            }
            
            success = self.cache_loader.save_cache(
                scenario_folder=scenario_folder,
                scenario_id=scenario_id,
                results=results,
                metadata=metadata
            )
            
            if success:
                self.stats['processed_scenarios'] += 1
                self.stats['total_components'] += len(results)
                logger.info(f"✅ 场景 {scenario_id} 处理完成: {len(results)} 个组件")
                return True, "处理成功", len(results)
            else:
                logger.error(f"❌ 场景 {scenario_id} 保存缓存失败")
                self.stats['failed_scenarios'] += 1
                return False, "保存缓存失败", 0
            
        except Exception as e:
            logger.error(f"❌ 场景 {scenario_id} 处理失败: {e}")
            self.stats['failed_scenarios'] += 1
            return False, f"处理失败: {str(e)}", 0
    
    def run(self, specific_scenarios: Optional[List[str]] = None):
        """
        运行预处理任务
        
        Args:
            specific_scenarios: 指定要处理的场景ID列表（None表示处理全部）
        """
        # 加载场景列表
        scenarios = self.load_scenarios(specific_scenarios)
        
        if not scenarios:
            logger.error("没有可处理的场景")
            return
        
        logger.info(f"\n{'='*70}")
        logger.info(f"开始预处理任务")
        logger.info(f"总场景数: {len(scenarios)}")
        logger.info(f"并发线程数: {self.workers}")
        logger.info(f"{'='*70}\n")
        
        # 使用线程池并发处理
        if self.workers > 1:
            # 多线程模式
            with ThreadPoolExecutor(max_workers=self.workers) as executor:
                # 提交任务
                future_to_scenario = {
                    executor.submit(self.process_scenario, scenario): scenario
                    for scenario in scenarios
                }
                
                # 收集结果
                for future in as_completed(future_to_scenario):
                    scenario = future_to_scenario[future]
                    try:
                        success, message, count = future.result()
                        if not success:
                            logger.error(f"场景 {scenario['scenario_id']} 失败: {message}")
                    except Exception as e:
                        logger.error(f"场景 {scenario['scenario_id']} 异常: {e}")
        else:
            # 单线程模式（便于调试）
            for scenario in scenarios:
                success, message, count = self.process_scenario(scenario)
                if not success:
                    logger.error(f"场景 {scenario['scenario_id']} 失败: {message}")
        
        # 打印统计信息
        self.print_statistics()
    
    def print_statistics(self):
        """打印统计信息"""
        end_time = datetime.now()
        duration = (end_time - self.stats['start_time']).total_seconds()
        
        logger.info(f"\n{'='*70}")
        logger.info(f"预处理任务完成")
        logger.info(f"{'='*70}")
        logger.info(f"总场景数: {self.stats['total_scenarios']}")
        logger.info(f"成功处理: {self.stats['processed_scenarios']}")
        logger.info(f"跳过（已有缓存）: {self.stats['skipped_scenarios']}")
        logger.info(f"处理失败: {self.stats['failed_scenarios']}")
        logger.info(f"总组件数: {self.stats['total_components']}")
        logger.info(f"总耗时: {duration:.2f} 秒")
        
        if self.stats['processed_scenarios'] > 0:
            avg_time = duration / self.stats['processed_scenarios']
            logger.info(f"平均每场景耗时: {avg_time:.2f} 秒")
        
        logger.info(f"{'='*70}\n")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='预处理指标数据，生成缓存文件',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
    # 处理所有场景（使用默认参数）
    python scripts/preprocess_metrics.py
    
    # 指定数据目录和并发数
    python scripts/preprocess_metrics.py --data-root /path/to/data --workers 8
    
    # 强制重新处理所有场景
    python scripts/preprocess_metrics.py --force-refresh
    
    # 仅处理指定场景
    python scripts/preprocess_metrics.py --scenarios "scene1,scene2,scene3"
        """
    )
    
    parser.add_argument(
        '--data-root',
        type=str,
        default='data/envs/aiops-22',
        help='数据集根目录（默认: data/envs/aiops-22）'
    )
    
    parser.add_argument(
        '--workers',
        type=int,
        default=4,
        help='并发线程数（默认: 4）'
    )
    
    parser.add_argument(
        '--force-refresh',
        action='store_true',
        help='强制重新处理已有缓存的场景'
    )
    
    parser.add_argument(
        '--scenarios',
        type=str,
        default=None,
        help='仅处理指定场景（逗号分隔，如 "scene1,scene2"）'
    )
    
    parser.add_argument(
        '--pre-fault-steps',
        type=int,
        default=2,
        help='MetricTool 的 pre_fault_steps 参数（默认: 2）'
    )
    
    parser.add_argument(
        '--contamination',
        type=float,
        default=0.1,
        help='MetricTool 的 contamination 参数（默认: 0.05）'
    )
    
    args = parser.parse_args()
    
    # 解析指定场景
    specific_scenarios = None
    if args.scenarios:
        specific_scenarios = [s.strip() for s in args.scenarios.split(',')]
        logger.info(f"仅处理指定场景: {specific_scenarios}")
    
    # 创建预处理器
    preprocessor = MetricPreprocessor(
        data_root=args.data_root,
        workers=args.workers,
        force_refresh=args.force_refresh,
        pre_fault_steps=args.pre_fault_steps,
        contamination=args.contamination
    )
    
    # 运行预处理任务
    try:
        preprocessor.run(specific_scenarios=specific_scenarios)
    except KeyboardInterrupt:
        logger.info("\n用户中断，退出预处理任务")
        sys.exit(1)
    except Exception as e:
        logger.error(f"预处理任务异常: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()

